package com.ai.module.mj.service.impl;

import com.ai.module.chat.common.constant.RedisKeyStringConstant;
import com.ai.module.chat.common.domain.vo.RiskReason;
import com.ai.module.chat.config.advice.exception.BusinessException;
import com.ai.module.chat.util.DateUtils;
import com.ai.module.chat.util.RedisTemplateUtils;
import com.ai.module.chat.util.aliyun.AliyunUtil;
import com.ai.common.BaseResp;
import com.ai.module.mj.constant.enums.BlendDimensions;
import com.ai.module.mj.constant.enums.TaskAction;
import com.ai.module.mj.constant.stringConstant.Constants;
import com.ai.module.mj.constant.stringConstant.ReturnCode;
import com.ai.module.mj.domain.Task;
import com.ai.module.mj.domain.bo.MJDataUrl;
import com.ai.module.mj.domain.vo.SubmitResultVO;
import com.ai.module.mj.loadbalancer.DiscordInstance;
import com.ai.module.mj.loadbalancer.DiscordLoadBalancer;
import com.ai.module.mj.service.HighTaskStoreService;
import com.ai.module.mj.service.LowTaskStoreService;
import com.ai.module.mj.service.TaskService;
import com.ai.module.mj.service.TaskStoreService;
import com.ai.module.mj.util.MimeTypeUtils;
import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.Nullable;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicReference;

import static com.ai.module.user.service.impl.Chat4ServiceImpl.ERROR_G_MESSAGE;

@Slf4j
@Service
public class TaskServiceImpl implements TaskService {

	public String ERROR_MESSAGE = ERROR_G_MESSAGE;
	@Autowired
	private HighTaskStoreService highTaskStoreService;

	@Autowired
	private LowTaskStoreService lowTaskStoreService;

	@Autowired
	private DiscordLoadBalancer discordLoadBalancer;

	private Boolean end = Boolean.FALSE;

	@Override
	public SubmitResultVO submitImagine(Task task) {
		// 进行 数据校验
		checkPrompt(task);
		DiscordInstance instance = this.discordLoadBalancer.chooseInstance();
		if (instance == null) {
			throw new BusinessException("任务排队人数多,请待会儿再试......");
		}
		task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, instance.getInstanceId());
		TaskStoreService taskStoreService = highTaskStoreService;
		if ( !Objects.equals(task.getDrawType(), "FAST")){
			taskStoreService = lowTaskStoreService;
		}
		if ( end ){
			taskStoreService.saveEndTask(task);
			return SubmitResultVO.of();
		}
		List<MJDataUrl> dataUrls = task.getDataUrls();
		List<String> fileUrlList = task.getFileUrlList();
		TaskStoreService finalTaskStoreService = taskStoreService;
		return instance.submitTask(task, () -> {
			List<String> imageUrls = new ArrayList<>();
			if (!CollectionUtils.isEmpty(fileUrlList)){
				imageUrls= fileUrlList;
			}else{
				BaseResp<Void> uploadResult = dealDataUrlList(task, dataUrls, instance, imageUrls);
				if (uploadResult != null) {
					return uploadResult;
				}
			}

			//  todo 图片必须限制 不然Redis 顶不住
			task.setDataUrls(new ArrayList<>());
			if (!imageUrls.isEmpty()) {
				String imageUrlListStr = String.join(" ", imageUrls);

				// 判断是不是 绘画 有没有参数标志
				String crefPrompt = task.getCrefPrompt();
				TaskAction action = task.getAction();
				if (StringUtils.hasText(crefPrompt)
						&& action == TaskAction.IMAGINE
				){
					String realCrefPrompt = String.format(crefPrompt, imageUrlListStr);
					task.setPromptEn(realCrefPrompt);

				}else{
					task.setRealPromptEn(imageUrlListStr + " " + task.getPromptEn());
				}
				task.setPrompt(task.getPrompt());
				task.setDescription("/imagine " + task.getPrompt());
				finalTaskStoreService.save(task);

			}

			// 真正执行的  画图
			return instance.imagine(task.getPromptEn(), task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE));
		});
	}

	@Nullable
	private static BaseResp<Void> dealDataUrlList(Task task, List<MJDataUrl> dataUrls, DiscordInstance instance, List<String> imageUrls) {
		for (MJDataUrl dataUrl : dataUrls) {
			String taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
			BaseResp<String> uploadResult = instance.upload(taskFileName, dataUrl);
			if (uploadResult.getCode() != ReturnCode.SUCCESS) {
				log.error("midjourney提交图片dataUrl失败[{}]", JSONObject.toJSONString(uploadResult));
				return BaseResp.of(uploadResult.getCode(), uploadResult.getMessage());
			}
			String finalFileName = uploadResult.getData();
			BaseResp<String> sendImageResult = instance.sendImageMessage("upload image: " + finalFileName, finalFileName);
			if (sendImageResult.getCode() != ReturnCode.SUCCESS) {
				log.error("midjourney提交图片dataUrl失败[{}]", JSONObject.toJSONString(sendImageResult));

				return BaseResp.of(sendImageResult.getCode(), sendImageResult.getMessage());
			}
			imageUrls.add(sendImageResult.getData());
		}
		return null;
	}

	private static void checkPrompt(Task task) {
		String prompt = task.getPrompt();
		if (!StringUtils.hasText(prompt)){
			return;
		}
		CompletableFuture<RiskReason> future = AliyunUtil.textModerationAsync(prompt);
		try {
			RiskReason riskReason = future.get();
			if (!riskReason.getSuccess()) {
				// 数据入库
				throw new BusinessException(ERROR_G_MESSAGE.replaceAll("@",""), ReturnCode.FAILURE);
			}
		} catch ( Exception e) {
			e.printStackTrace();
			log.error("文本审核异常");
			if (e instanceof BusinessException){
				throw (BusinessException) e;
			}
		}
	}

	@Override
	public SubmitResultVO submitUpscale(Task task, String targetMessageId, String targetMessageHash, int index, int messageFlags) {

		String instanceId = task.getPropertyGeneric(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID);
		DiscordInstance discordInstance = this.discordLoadBalancer.getDiscordInstance(instanceId);
		if (discordInstance == null || !discordInstance.isAlive()) {
			throw new BusinessException(ReturnCode.NOT_FOUND, "账号不可用: " + instanceId);
		}
		return discordInstance.submitTask(task, () -> discordInstance.upscale(targetMessageId, index, targetMessageHash, messageFlags, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE)));
	}

	@Override
	public SubmitResultVO submitVariation(Task task, String targetMessageId, String targetMessageHash, int index, int messageFlags) {
		String instanceId = task.getPropertyGeneric(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID);
		DiscordInstance discordInstance = this.discordLoadBalancer.getDiscordInstance(instanceId);
		if (discordInstance == null || !discordInstance.isAlive()) {
			throw new BusinessException(ReturnCode.NOT_FOUND, "账号不可用: " + instanceId);
		}
		return discordInstance.submitTask(task, () -> discordInstance.variation(targetMessageId, index, targetMessageHash, messageFlags, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE)));
	}

	@Override
	public SubmitResultVO submitReroll(Task task, String targetMessageId, String targetMessageHash, int messageFlags) {
		String instanceId = task.getPropertyGeneric(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID);
		DiscordInstance discordInstance = this.discordLoadBalancer.getDiscordInstance(instanceId);
		if (discordInstance == null || !discordInstance.isAlive()) {
			throw new BusinessException(ReturnCode.NOT_FOUND, "账号不可用: " + instanceId);
		}
		return discordInstance.submitTask(task, () -> discordInstance.reroll(targetMessageId, targetMessageHash, messageFlags, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE)));
	}

	@Override
	public SubmitResultVO submitDescribe(Task task,
										 MJDataUrl dataUrl,
										 List<String> fileUrlList) {
		DiscordInstance discordInstance = this.discordLoadBalancer.chooseInstance();
		if (discordInstance == null) {
			throw new BusinessException(ReturnCode.FAILURE, "无可用的账号实例");
		}
		task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, discordInstance.getInstanceId());
		return discordInstance.submitTask(task, () -> {
			if (CollectionUtils.isEmpty(fileUrlList)){
				String taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
				BaseResp<String> uploadResult = discordInstance.upload(taskFileName, dataUrl);
				if (uploadResult.getCode() != ReturnCode.SUCCESS) {
					log.error("midjourney执行任务失败[{}]", JSONObject.toJSONString(uploadResult));

					return BaseResp.of(uploadResult.getCode(), uploadResult.getMessage());
				}
				String finalFileName = uploadResult.getData();
				return discordInstance.describe(finalFileName, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE));

			}else {
				RedisTemplateUtils.set(String.format(RedisKeyStringConstant.DESC_URL_TASK_ID, fileUrlList.get(0)),task.getId(), DateUtils.TIME_DAY    );
				return discordInstance.describe(fileUrlList.get(0), task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE));

			}
		});
	}

	@Override
	public SubmitResultVO submitBlend(Task task, List<MJDataUrl> dataUrls, BlendDimensions dimensions) {
		checkPrompt(task);
		DiscordInstance discordInstance = this.discordLoadBalancer.chooseInstance();
		if (discordInstance == null) {
			throw new BusinessException(ReturnCode.NOT_FOUND, "无可用的账号实例");
		}
		task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, discordInstance.getInstanceId());
		AtomicReference<List<String>> fileUrlList = new AtomicReference<>(task.getFileUrlList());
		return discordInstance.submitTask(task, () -> {
//			List<String> finalFileNames = fileUrlList;

			if (!CollectionUtils.isEmpty(dataUrls)){
				fileUrlList.set(new ArrayList<>());
				for (MJDataUrl dataUrl : dataUrls) {
					String taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
					BaseResp<String> uploadResult = discordInstance.upload(taskFileName, dataUrl);
					if (uploadResult.getCode() != ReturnCode.SUCCESS) {
						log.error("midjourney执行拼图上传失败[{}]", JSONObject.toJSONString(uploadResult));
						throw new BusinessException("系统繁忙", ReturnCode.FAILURE);
					}
					fileUrlList.get().add(uploadResult.getData());
				}
			}

			return discordInstance.blend(fileUrlList.get(), dimensions, task.getPropertyGeneric(Constants.TASK_PROPERTY_NONCE));
		});
	}

	@Override
	public void setEnd() {
		this.end = true;
	}
}
